package com.boge.ai.rag;

import com.boge.ai.entity.milvus.MilvusEntity;
import com.boge.ai.entity.rag.FaqItem;
import com.boge.ai.utils.LLMUtils;
import com.boge.ai.utils.MilvusUtils;
import io.milvus.response.QueryResultsWrapper;
import io.milvus.response.SearchResultsWrapper;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

public class Demo06 {
    public static void main(String[] args) throws Exception{
        MilvusUtils mu = new MilvusUtils();
        // 1. 提示词
        String prompt = "得了白癜风怎么办?";
        // 2. 需要对提示词通过相同的向量模版来向量化
        float[] embedding = LLMUtils.embedding(prompt, null);
        // 3.根据获取到的向量值去 Milvus 中匹配对应的 数据
        List<Float> floats = new ArrayList<>();
        for (float v : embedding) {
            floats.add(v);
        }
        // 获取到对应结果的 包装对象
        SearchResultsWrapper searchResultsWrapper = mu.searchVectorData(floats);
        List<QueryResultsWrapper.RowRecord> rowRecords = searchResultsWrapper.getRowRecords();
        // 通过集合来记录 从向量数据库中 检索到的匹配的信息
        List<FaqItem> faqItems = new ArrayList<>();
        if(rowRecords != null && rowRecords.size() > 0){
            for (QueryResultsWrapper.RowRecord rowRecord : rowRecords) {
                Map<String, Object> fieldValues = rowRecord.getFieldValues();
                // 原始的问题
                String instruction = fieldValues.get(MilvusEntity.Field.INSTRUCTION).toString();
                // 获取问题对应的答案
                String output = fieldValues.get(MilvusEntity.Field.OUTPUT).toString();

                FaqItem faqItem = new FaqItem();
                faqItem.setInstruction(instruction);
                faqItem.setOutput(output);
                faqItems.add(faqItem);
            }
        }

        System.out.println(faqItems);
        System.out.println("--------------------------");
        // 处理下匹配到的信息
        String result = faqItems.stream()
                .map(item-> "Q:"+ item.getInstruction() +"\n A:" + item.getOutput())
                .collect(Collectors.joining("\n"));
        // 4. 原始的提示词 + 相关的提示信息 组合成新的提示词
        prompt = String.format("""
                # 角色
                你是一个问答机器人。
                你的任务是根据下述给定的已知信息回答用户问题。
                根据已知信息的内容推理给出用户问题的解决方案
                确保你的回复完全依据下述已知信息。不要编造答案。
                如果下述已知信息不足以回答用户的问题，请直接回复"我无法回答您的问题"。
                
                已知信息:
                %s
                
                用户问题：
                %s
                
                请用中文回答用户问题。
                """, result, prompt);
        String completion = LLMUtils.completion(prompt, null);
        System.out.println(completion);
    }
}
